import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import matplotlib.animation as animation
from IPython.display import HTML
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import plotly.express as px
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()
def get_stock_data(ts_code, start_date, end_date):
    """从Tushare获取股票数据"""
    df = ts.pro_bar(ts_code=ts_code, adj='qfq', start_date=start_date, end_date=end_date)
    df = df.sort_values('trade_date')
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df.set_index('trade_date', inplace=True)
    return df


def double_ma_strategy(data, short_window=5, long_window=20):
    """双均线策略实现"""
    signals = pd.DataFrame(index=data.index)
    signals['price'] = data['close']
    signals['short_ma'] = data['close'].rolling(window=short_window, min_periods=1).mean()
    signals['long_ma'] = data['close'].rolling(window=long_window, min_periods=1).mean()
    signals['signal'] = 0.0
    signals['signal'][short_window:] = np.where(
        signals['short_ma'][short_window:] > signals['long_ma'][short_window:], 1.0, 0.0)
    signals['positions'] = signals['signal'].diff()
    return signals


def backtest_strategy(signals, initial_capital=100000.0):
    """回测策略"""
    positions = pd.DataFrame(index=signals.index).fillna(0.0)
    positions['stock'] = 100 * signals['signal']  # 假设每次交易100股
    portfolio = positions.multiply(signals['price'], axis=0)
    pos_diff = positions.diff()

    portfolio['holdings'] = (positions.multiply(signals['price'], axis=0)).sum(axis=1)
    portfolio['cash'] = initial_capital - (pos_diff.multiply(signals['price'], axis=0)).sum(axis=1).cumsum()
    portfolio['total'] = portfolio['cash'] + portfolio['holdings']
    portfolio['returns'] = portfolio['total'].pct_change()
    return portfolio


# 获取数据并测试双均线策略
data = get_stock_data('600519.SH', '20220101', '2025331')
signals = double_ma_strategy(data)
portfolio = backtest_strategy(signals)
# 计算策略表现
returns = portfolio['returns']
cumulative_returns = (1 + returns).cumprod() - 1
annual_return = returns.mean() * 252
volatility = returns.std() * np.sqrt(252)
sharpe_ratio = annual_return / volatility

print(f"双均线策略表现:")
print(f"累计收益率: {cumulative_returns[-1]:.2%}")
print(f"年化收益率: {annual_return:.2%}")
print(f"年化波动率: {volatility:.2%}")
print(f"夏普比率: {sharpe_ratio:.2f}")
# 绘制结果
def create_interactive_plot(signals, portfolio):
    """使用Plotly创建交互式图表"""

    # 创建带两个子图的图表
    fig = make_subplots(rows=2, cols=1, shared_xaxes=True,
                        vertical_spacing=0.05,
                        subplot_titles=('Price and Moving Averages', 'Portfolio Value'))

    # 第一个子图：价格和均线
    fig.add_trace(
        go.Scatter(x=signals.index, y=signals['price'],
                   name='Price', line=dict(color='black', width=1)),
        row=1, col=1
    )

    fig.add_trace(
        go.Scatter(x=signals.index, y=signals['short_ma'],
                   name='Short MA', line=dict(width=2)),
        row=1, col=1
    )

    fig.add_trace(
        go.Scatter(x=signals.index, y=signals['long_ma'],
                   name='Long MA', line=dict(width=2)),
        row=1, col=1
    )

    # 添加买入信号标记
    buy_signals = signals[signals.positions == 1.0]
    fig.add_trace(
        go.Scatter(x=buy_signals.index, y=buy_signals['short_ma'],
                   mode='markers', name='Buy',
                   marker=dict(symbol='triangle-up', size=10, color='green')),
        row=1, col=1
    )

    # 添加卖出信号标记
    sell_signals = signals[signals.positions == -1.0]
    fig.add_trace(
        go.Scatter(x=sell_signals.index, y=sell_signals['short_ma'],
                   mode='markers', name='Sell',
                   marker=dict(symbol='triangle-down', size=10, color='red')),
        row=1, col=1
    )

    # 第二个子图：投资组合价值
    fig.add_trace(
        go.Scatter(x=portfolio.index, y=portfolio['total'],
                   name='Portfolio Value', line=dict(width=2)),
        row=2, col=1
    )

    # 更新图表布局
    fig.update_layout(
        height=800,
        title_text="Trading Strategy Performance",
        hovermode="x unified",
        showlegend=True
    )

    # 更新y轴标签
    fig.update_yaxes(title_text="Price (RMB)", row=1, col=1)
    fig.update_yaxes(title_text="Portfolio Value (RMB)", row=2, col=1)

    # 添加范围滑块
    fig.update_layout(
        xaxis=dict(
            rangeselector=dict(
                buttons=list([
                    dict(count=1, label="1m", step="month", stepmode="backward"),
                    dict(count=6, label="6m", step="month", stepmode="backward"),
                    dict(count=1, label="YTD", step="year", stepmode="todate"),
                    dict(count=1, label="1y", step="year", stepmode="backward"),
                    dict(step="all")
                ])
            ),
            rangeslider=dict(visible=True),
            type="date"
        )
    )
    # 在create_interactive_plot函数中，可以添加这些定制选项：

    # 添加交易信号注释
    for date, row in buy_signals.iterrows():
        fig.add_annotation(x=date, y=row['short_ma'],
                           text="Buy", showarrow=True,
                           arrowhead=1, ax=0, ay=-30,
                           row=1, col=1)

    # 添加投资组合关键点标记
    max_value = portfolio['total'].max()
    max_date = portfolio['total'].idxmax()
    fig.add_annotation(x=max_date, y=max_value,
                       text=f"Max: {max_value:,.0f}",
                       showarrow=True, arrowhead=1,
                       row=2, col=1)

    # 更改主题
    fig.update_layout(template="plotly_dark")

    # 添加自定义悬停文本
    fig.update_traces(
        hovertemplate="<b>%{x|%Y-%m-%d}</b><br>Value: %{y:,.2f} RMB"
    )
    return fig
# 使用示例
plotly_fig = create_interactive_plot(signals, portfolio)
plotly_fig.show()

def prepare_features(data, window=5):
    """准备机器学习特征"""
    df = data.copy()

    # 创建技术指标特征
    df['returns'] = df['close'].pct_change()
    df['volatility'] = df['returns'].rolling(window).std()
    df['momentum'] = df['close'] / df['close'].shift(window) - 1
    df['ma5'] = df['close'].rolling(5).mean()
    df['ma20'] = df['close'].rolling(20).mean()
    df['ma_diff'] = df['ma5'] - df['ma20']
    df['high_5'] = df['high'].rolling(window).max()
    df['low_5'] = df['low'].rolling(window).min()

    # 创建目标变量 (未来5天收益率是否为正)
    df['target'] = np.where(df['close'].shift(-5) > df['close'], 1, 0)

    # 删除缺失值
    df.dropna(inplace=True)

    return df


def train_ml_model(features):
    """训练机器学习模型"""
    X = features[['returns', 'volatility', 'momentum', 'ma_diff', 'high_5', 'low_5']]
    y = features['target']

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)

    # 训练随机森林模型
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)

    # 评估模型
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print(f"模型准确率: {accuracy:.2%}")

    return model


def prepare_features(data, window=5):
    """准备机器学习特征"""
    df = data.copy()

    # 创建技术指标特征
    df['returns'] = df['close'].pct_change()
    df['volatility'] = df['returns'].rolling(window).std()
    df['momentum'] = df['close'] / df['close'].shift(window) - 1
    df['ma5'] = df['close'].rolling(5).mean()
    df['ma20'] = df['close'].rolling(20).mean()
    df['ma_diff'] = df['ma5'] - df['ma20']
    df['high_5'] = df['high'].rolling(window).max()
    df['low_5'] = df['low'].rolling(window).min()

    # 创建目标变量 (未来5天收益率是否为正)
    df['target'] = np.where(df['close'].shift(-5) > df['close'], 1, 0)

    # 删除缺失值
    df.dropna(inplace=True)

    return df


def train_ml_model(features):
    """训练机器学习模型"""
    X = features[['returns', 'volatility', 'momentum', 'ma_diff', 'high_5', 'low_5']]
    y = features['target']

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)

    # 训练随机森林模型
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)

    # 评估模型
    y_pred = model.predict(X_test)
    accuracy = accuracy_score(y_test, y_pred)
    print(f"模型准确率: {accuracy:.2%}")

    return model


def ml_strategy(data, model):
    """基于机器学习的交易策略"""
    features = prepare_features(data)
    X = features[['returns', 'volatility', 'momentum', 'ma_diff', 'high_5', 'low_5']]

    # 确保我们只对有效的特征数据进行预测
    valid_dates = features.index
    predictions = model.predict(X)

    # 创建信号DataFrame，只包含有预测结果的日期
    signals = pd.DataFrame(index=valid_dates)
    signals['price'] = data.loc[valid_dates, 'close']
    signals['signal'] = predictions
    signals['positions'] = signals['signal'].diff()

    return signals
# 准备数据并训练模型
full_data = get_stock_data('600519.SH', '20220101', '2025231')
features = prepare_features(full_data)
model = train_ml_model(features)
# 应用机器学习策略
ml_signals = ml_strategy(full_data, model)
ml_portfolio = backtest_strategy(ml_signals)
# 绘制结果
fig = plt.figure(figsize=(12, 8))
ax1 = fig.add_subplot(211, ylabel='Price in RMB')
ml_signals['price'].plot(ax=ax1, color='k', lw=1.)

ax1.plot(ml_signals.loc[ml_signals.positions == 1.0].index,
         ml_signals.price[ml_signals.positions == 1.0],
         '^', markersize=10, color='g')
ax1.plot(ml_signals.loc[ml_signals.positions == -1.0].index,
         ml_signals.price[ml_signals.positions == -1.0],
         'v', markersize=10, color='r')
ax2 = fig.add_subplot(212, ylabel='Portfolio value in RMB')
ml_portfolio['total'].plot(ax=ax2, lw=2.)
plt.show()
# 计算策略表现
ml_returns = ml_portfolio['returns']
ml_cumulative_returns = (1 + ml_returns).cumprod() - 1
ml_annual_return = ml_returns.mean() * 252
ml_volatility = ml_returns.std() * np.sqrt(252)
ml_sharpe_ratio = ml_annual_return / ml_volatility

print(f"\n机器学习策略表现:")
print(f"累计收益率: {ml_cumulative_returns[-1]:.2%}")
print(f"年化收益率: {ml_annual_return:.2%}")
print(f"年化波动率: {ml_volatility:.2%}")
print(f"夏普比率: {ml_sharpe_ratio:.2f}")

# 比较两种策略
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
comparison = pd.DataFrame({
    '双均线策略': [cumulative_returns[-1], annual_return, volatility, sharpe_ratio],
    '机器学习策略': [ml_cumulative_returns[-1], ml_annual_return, ml_volatility, ml_sharpe_ratio]
}, index=['累计收益率', '年化收益率', '年化波动率', '夏普比率'])

print("\n策略比较:")
print(comparison)


# 绘制累计收益率比较
plt.figure(figsize=(12, 6))
plt.plot((1 + returns).cumprod() - 1, label='双均线策略')
plt.plot((1 + ml_returns).cumprod() - 1, label='机器学习策略')
plt.title('策略累计收益率比较')
plt.xlabel('日期')
plt.ylabel('累计收益率')
plt.legend()
plt.grid(True)
plt.show()